package com.it.state;

import com.it.operator.utils.SourceUtils;
import com.it.pojo.Event;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

import java.util.ArrayList;
import java.util.List;

/**
 * 批量写出以及故障恢复.
 * @author code1997
 */
public class BufferingSinkDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();
        SingleOutputStreamOperator<Event> eventSource = SourceUtils.getEventSource(executionEnvironment);
        eventSource.print("input");
        eventSource.addSink(new BufferingSink(5));
        executionEnvironment.execute();
    }

    public static class BufferingSink implements SinkFunction<Event>, CheckpointedFunction {

        private final int bufferSize;
        private final List<Event> bufferedElements;
        private ListState<Event> checkPointedState;

        public BufferingSink(int bz) {
            if (bz < 1) {
                throw new IllegalArgumentException("buffer size must be more than 1");
            }
            this.bufferSize = bz;
            this.bufferedElements = new ArrayList<>();
        }

        /**
         * 每来一个数据我们需要做的.
         */
        @Override
        public void invoke(Event value, Context context) throws Exception {
            bufferedElements.add(value);
            //判断如果达到阈值，就批量写入
            if (bufferedElements.size() == bufferSize) {
                //写入数据到外部系统
                for (Event event : bufferedElements) {
                    System.out.println(event);
                }
                System.out.println("========写出完毕=======");
                bufferedElements.clear();

            }
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            //需要清理状态
            checkPointedState.clear();
            //对状态进行持久化，复制缓存的列表到列表状态
            for (Event event : bufferedElements) {
                checkPointedState.add(event);
            }
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            //初始化状态
            ListStateDescriptor<Event> listStateDescriptor = new ListStateDescriptor<>("buffered_states", Types.POJO(Event.class));
            checkPointedState = context.getOperatorStateStore().getListState(listStateDescriptor);
            //如果从故障恢复，需要将ListState中所有元素复制到列表中
            if (context.isRestored()) {
                for (Event event : checkPointedState.get()) {
                    bufferedElements.add(event);
                }
            }
        }
    }
}
